'''用于绘制陨石坑领域密度图像的脚本'''
import numpy as np
import yaml
from reader import from_TIFF, from_IMG

with open(
    "/home/a804_cbf/Code/luner_crater/simulator/labeler/utils/craters.csv", "r"
) as f:
    craters = []
    for line in f:
        lat, lon, r = map(float, line.split(","))
        craters.append((lat, lon, r))
    craters = np.array(craters)

config_path = "/home/a804_cbf/Code/luner_crater/config/2080x4/HG-HV.yaml"
with open(config_path, "r") as f:
    config = yaml.load(f, Loader=yaml.FullLoader)
    nac_config = config["nac"]
    dem_config = config["dem"]
    subfig_config = config["subfig"]
dem_prj = from_TIFF(**dem_config, **subfig_config)
nac_prj = from_IMG(nac_config, **subfig_config)
nac = nac_prj.image.copy()
dem = dem_prj.image.copy()
x, y = nac_prj.forward(craters[:, 0], craters[:, 1])
r = craters[:, 2]
r_th = 100
rate = 5
with open("neighbor.csv", "w") as w:
    for xi, yi, ri in zip(x, y, r):
        neighbor = (
            ((x - xi) ** 2 + (y - yi) ** 2 < r_th**2) & (r / ri < (1 + rate / 100))
        ).sum() - 1
        print(neighbor, file=w)
